package org.taj.adapter.server;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.taj.adapter.llama.LLaMAClient;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

@Component
public class Server {
    private ZContext ctx;
    @Value("${adapter.req-port}")
    private int reqPort;
    @Value("${adapter.pub-port}")
    private int pubPort;
    @Autowired
    private LLaMAClient client;

    private AtomicLong sessionId = new AtomicLong(0);

    private Map<String, Session> sessionHolder = new HashMap<>();

    @PostConstruct
    public void start() {
        ctx = new ZContext();
        ZMQ.Socket repSocket = ctx.createSocket(SocketType.REP);
        repSocket.bind("tcp://*:" + reqPort);
        ZMQ.Socket pubSocket = ctx.createSocket(SocketType.PUB);
        pubSocket.bind("tcp://*:" + pubPort);
        while (true) {
            String header = repSocket.recvStr();
            String body = repSocket.recvStr();
            JSONObject jsonBody = JSON.parseObject(body);
            String systemPrompt = jsonBody.getString("systemPrompt");
            String userInput = jsonBody.getString("userInput");
            Session session;
            if(header.equals("00000000")) {
                do {
                    header = String.format("%08d", sessionId.addAndGet(1));
                } while (sessionHolder.containsKey(header));
                session = new Session();
                sessionHolder.put(header, session);
            }
            else if(sessionHolder.containsKey(header)) {
                session = sessionHolder.get(header);
            } else {
                session = new Session();
                sessionHolder.put(header, session);
            }
            repSocket.send(header, ZMQ.DONTWAIT);
            session.setPrompt(systemPrompt);
            session.addUserInput(userInput);
            client.send(session, pubSocket, header);


        }
    }
}
